import os
import pathlib
import unittest

import cv2
import numpy as np
from matplotlib import pyplot as plt
from ocr_process.infer.predict import TextImagePredict, UvdocPredict, ClasPredict

from taskinfo.Uie import Uie


class ImageTest(unittest.TestCase):


    def test_cv3(self):
        # 读取图像
        image = cv2.imread('4.jpg')
        # 将蓝色通道设置为0
        image[:, :, 0] = 0  # OpenCV的通道顺序是BGR，第1个通道（索引为0）是蓝色通道
        image[:, :, 1] = 0  # OpenCV的通道顺序是BGR，第1个通道（索引为0）是蓝色通道

        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 显示或保存处理后的图像
        cv2.imshow("Image without Blue Channel", image)
        cv2.imshow("gray_image", gray_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def test_enhance_image(self):
        image = cv2.imread('r.jpg')
        # 增强对比度，减小亮度
        alpha = 1.8  # 对比度增益因子，值越大对比越强
        beta = -55  # 亮度调低值，负值会使黑色字体更突出

        # 调整对比度和亮度
        enhanced_image = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)

        # 显示增强后的图像
        cv2.imshow("Enhanced Black Font", enhanced_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        cv2.imwrite("r_enhanced_image.jpg", enhanced_image)

    def test_cv2(self):
        image = cv2.imread('4.jpg')
        _, _, r = cv2.split(image)
        cv2.imwrite("r.jpg", r)


    def test_rename(self):
        path_dir = pathlib.Path(r'D:\ocr\train_data\1114_rongcheng\red_channel')
        for file in path_dir.iterdir():
            file_name = file.name
            new_file_name = '1114_' + file_name
            file.rename(path_dir / new_file_name)





    def test_uvdoc_file(self):
        uvdoc_predict = UvdocPredict(os.path.join("D:/ocr/onnx/deploy/uvdoc", "inference.onnx"))
        file_dir_path = pathlib.Path(r"D:\ocr\train_data\1111_rongcheng\TextImagePredict")
        # 过滤图片文件
        image_files = [f for f in file_dir_path.iterdir() if (f.suffix == '.png' or f.suffix == '.jpg')]
        for image_file in image_files:
            try:
                image = cv2.imread(str(image_file))
                input = np.array(image)
                image_np = uvdoc_predict(input)
                cv2.imwrite(rf'D:\ocr\train_data\1111_rongcheng\uvdoc\{image_file.name}', image_np)
                print("----------")
            except (Exception):
                print("图片转换异常")
                continue

    def test_r_file(self):
        file_dir_path = pathlib.Path(r"D:\ocr\train_data\1114_rongcheng\TextImagePredict")
        # 过滤图片文件
        image_files = [f for f in file_dir_path.iterdir() if (f.suffix == '.png' or f.suffix == '.jpg')]
        for image_file in image_files:
            try:
                image = cv2.imread(str(image_file))
                _, _, R = cv2.split(image)
                cv2.imwrite(rf'D:\ocr\train_data\1114_rongcheng\red_channel\{image_file.name}', R)
                print("----------")
            except (Exception):
                print("图片转换异常")
                continue

    def test_list_file(self):
        text_image_predict = TextImagePredict(os.path.join("D:/ocr/onnx/deploy/text_image", "inference.onnx"))
        file_dir_path = pathlib.Path(r"D:\ocr\train_data\1114_rongcheng")
        # 过滤图片文件
        image_files = [f for f in file_dir_path.iterdir() if (f.suffix == '.png' or f.suffix == '.jpg')]
        for image_file in image_files:

            try:
                image = cv2.imread(str(image_file))
                input = np.array(image)
                image_np = text_image_predict(input)
                cv2.imwrite(rf'D:\ocr\train_data\1114_rongcheng\TextImagePredict\{image_file.name}', image_np)
            except (Exception):
                print("图片转换异常")
                continue



    def test_image_predict(self):
        # 读取图像
        image = cv2.imread('1.jpg')
        input = np.array(image)
        text_image_predict = TextImagePredict(os.path.join("", "inference.onnx"))
        image_np = text_image_predict(input)
        cv2.imwrite('text_mask.png', image_np)




    def test_image_mask4(self):
        # 读取图像
        image = cv2.imread('1.jpg')

        # 转换到HSV颜色空间
        hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

        # 定义红色的HSV范围
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([150, 100, 100])
        upper_red2 = np.array([190, 255, 255])

        # 定义黄色的HSV范围
        lower_yellow = np.array([20, 100, 100])
        upper_yellow = np.array([30, 255, 255])

        # 创建红色掩码
        mask_red1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
        mask_red = mask_red1 | mask_red2

        # 创建黄色掩码
        mask_yellow = cv2.inRange(hsv_image, lower_yellow, upper_yellow)

        # 合并掩码
        combined_mask = mask_red | mask_yellow

        # 反转掩码，去除红色和黄色
        final_mask = cv2.bitwise_not(combined_mask)

        # 应用掩码到原始图像
        result_image = cv2.bitwise_and(image, image, mask=final_mask)

        # 显示结果
        cv2.imshow('Original Image', image)
        cv2.imshow('Result Image', result_image)

        # 保存结果
        # cv2.imwrite('result_image.jpg', result_image)

        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def test_image_mask3(self):
        # Load the new image
        image_path = 'logo2.png'
        image = cv2.imread(image_path)

        # Convert to grayscale to emphasize text contrast
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # Apply thresholding to create a binary mask
        _, mask = cv2.threshold(gray_image, 150, 255, cv2.THRESH_BINARY_INV)

        # Display the generated mask
        plt.imshow(mask, cmap='gray')
        plt.axis('off')
        plt.title("Generated Mask for Text")
        plt.show()

        # Save the mask
        cv2.imwrite('text_mask.png', mask)

    def test_image_mask2(self):
        image_path = 'logo.png'
        image = cv2.imread(image_path)
        image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        # Define HSV color range for yellow and red (adjusted to match the watermark colors)

        # Yellow color range
        lower_yellow = np.array([20, 100, 100])
        upper_yellow = np.array([30, 255, 255])

        # Red color range (OpenCV splits red at the beginning and end of the hue scale)
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([160, 100, 100])
        upper_red2 = np.array([180, 255, 255])

        # Create masks for yellow and red colors
        mask_yellow = cv2.inRange(image_hsv, lower_yellow, upper_yellow)
        mask_red1 = cv2.inRange(image_hsv, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(image_hsv, lower_red2, upper_red2)

        # Combine red masks and then combine with yellow mask
        mask_red = cv2.bitwise_or(mask_red1, mask_red2)
        mask = cv2.bitwise_or(mask_yellow, mask_red)

        # Display the generated mask
        plt.imshow(mask, cmap='gray')
        plt.axis('off')
        plt.title("Generated Mask")
        plt.show()

        # Save the mask
        # cv2.imwrite('/mnt/data/generated_mask.png', mask)

    def test_image_mask(self):
        # Load the image
        image_path = 'logo.png'
        image = cv2.imread(image_path)

        # Convert to RGB for displaying purposes in matplotlib
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # Display the original image
        plt.imshow(image_rgb)
        plt.axis('off')
        plt.title("Original Image")
        plt.show()

        # Convert to HSV color space for color analysis
        image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

        # Display the HSV channels separately to examine color ranges
        hue, saturation, value = cv2.split(image_hsv)

        # Plot HSV channels
        fig, ax = plt.subplots(1, 3, figsize=(15, 5))
        ax[0].imshow(hue, cmap='gray')
        ax[0].set_title("Hue Channel")
        ax[1].imshow(saturation, cmap='gray')
        ax[1].set_title("Saturation Channel")
        ax[2].imshow(value, cmap='gray')
        ax[2].set_title("Value Channel")

        for a in ax:
            a.axis('off')

        plt.show()

    def test_image_split(self):
        # cv2 通道分层
        image = cv2.imread("1.jpg")
        # BGR分层处理
        b, g, r = cv2.split(image)
        # cv2.imshow("blue", b)
        # cv2.imshow("green", g)
        # cv2.imshow("red", r)
        # 保持图片
        cv2.imwrite("red.jpg", r)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def test_image_resize(self):
        uie = Uie(schema_list=["车号", "供货单位", "收货单位", "物资品名", "毛重", "皮重", "净重", "对方重量", "船号"],
                  onnx_home="D:/ocr/onnx/deploy/rong_cheng", text_system_path="D:/ocr/onnx/deploy/text_system")

        res = uie("1.jpg")
        # [{'车号': [{'text': '车号：津A01292F', 'start': 22, 'end': 33, 'probability': 0.7960435440651565, 'bbox': [[338, 437, 740, 512]]}],
        #   '供货单位': [
        #       {'text': '天津中联进出口贸易有限公司', 'start': 43, 'end': 56, 'probability': 0.9978168655221964, 'bbox': [[519, 550, 1069, 621]]}],
        #   '收货单位': [{'text': '烧结库房|组织：/', 'start': 60, 'end': 69, 'probability': 0.9899582507847455, 'bbox': [[508, 659, 860, 710]]}],
        #   '物资品名': [{'text': '铁矿粉-混合粉', 'start': 72, 'end': 79, 'probability': 0.9958183206942266, 'bbox': [[495, 766, 795, 837]]}],
        #   '毛重': [{'text': '48.26吨202/11/47:29:03', 'start': 81, 'end': 102, 'probability': 0.9964236008586198, 'bbox': [[421, 890, 1015, 954]]}],
        #   '皮重': [
        #       {'text': '19.66吨2024/11/48:15:24', 'start': 104, 'end': 126, 'probability': 0.9975494828564848, 'bbox': [[414, 1003, 1021, 1073]]}],
        #   '净重': [{'text': '28.60吨', 'start': 128, 'end': 134, 'probability': 0.980724286308714, 'bbox': [[379, 1115, 594, 1197]]}],
        #   '对方重量': [{'text': '28.76吨扣0.0000', 'start': 146, 'end': 159, 'probability': 0.9919858066951406, 'bbox': [[491, 1386, 993, 1455]]}],
        #   '船号': [{'text': '比速6', 'start': 139, 'end': 142, 'probability': 0.86436614317725, 'bbox': [[464, 1236, 647, 1343]]}]}]
        print(res)
